package madixin.t637;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

class Solution {

    //DFS
    public List<Double> averageOfLevels(TreeNode root) {
        List<Integer> count = new ArrayList<>();
        List<Double> res = new ArrayList<>();

        average(root, 0, count, res);

        for (int i = 0; i < count.size(); i++) {
            res.set(i, res.get(i) / count.get(i));
        }
        return res;
    }

    private void average(TreeNode node, int i, List<Integer> count, List<Double> res) {
        if (node == null) {
            return;
        }

        if (i < count.size()) {
            count.set(i, count.get(i) + 1);
            res.set(i, res.get(i) + node.val);
        } else {
            count.add(1);
            res.add((double) node.val);
        }

        if (node.left != null) {
            average(node.left, i + 1, count, res);
        }

        if (node.right != null) {
            average(node.right, i + 1, count, res);
        }
    }


    //BST 广度优先算法
    //使用 BFS 进行层次遍历。不需要使用两个队列来分别存储当前层的节点和下一层的节点，
    //因为在开始遍历一层的节点时，当前队列中的节点数就是当前层的节点数，只要控制遍历这么多节点数，就能保证这次遍历的都是当前层的节点。
    public List<Double> averageOfLevelsBST(TreeNode root) {
        List<Double> res = new LinkedList<>();
        if (root == null) {
            return res;
        }

        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(root);
        while (!queue.isEmpty()) {
            int count = queue.size();
            double sum = 0;
            for (int i = 0; i < count; i++) {
                TreeNode curNode = queue.poll();
                sum += curNode.val;
                if (curNode.left != null) {
                    queue.add(curNode.left);
                }
                if (curNode.right != null) {
                    queue.add(curNode.right);
                }
            }
            res.add(sum / count);
        }
        return res;
    }
}